#!/usr/bin/env python3
"""
简化版 Prometheus 服务器
提供基本的指标收集、存储和查询功能
"""

import time
import json
import threading
import requests
import re
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
from datetime import datetime, timedelta
import yaml

class PrometheusServer:
    def __init__(self, config_file='prometheus.yml'):
        self.config = self.load_config(config_file)
        self.metrics_storage = {}  # 简单的内存存储
        self.targets = []
        self.scrape_interval = 15
        self.load_targets()
        
    def load_config(self, config_file):
        """加载配置文件"""
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except Exception as e:
            print(f"❌ 加载配置文件失败: {e}")
            return {}
    
    def load_targets(self):
        """从配置文件加载监控目标"""
        if 'scrape_configs' in self.config:
            for job in self.config['scrape_configs']:
                job_name = job.get('job_name', 'unknown')
                scrape_interval = job.get('scrape_interval', '15s')
                
                if 'static_configs' in job:
                    for static_config in job['static_configs']:
                        for target in static_config.get('targets', []):
                            self.targets.append({
                                'job': job_name,
                                'target': target,
                                'scrape_interval': scrape_interval
                            })
        
        print(f"📋 加载了 {len(self.targets)} 个监控目标:")
        for target in self.targets:
            print(f"  - {target['job']}: {target['target']}")
    
    def scrape_target(self, target_info):
        """抓取单个目标的指标"""
        target = target_info['target']
        job = target_info['job']
        
        try:
            url = f"http://{target}/metrics"
            response = requests.get(url, timeout=10)
            
            if response.status_code == 200:
                timestamp = int(time.time())
                metrics_text = response.text
                
                # 解析指标数据
                metrics = self.parse_metrics(metrics_text, job, target)
                
                # 存储到内存中
                for metric_name, metric_data in metrics.items():
                    if metric_name not in self.metrics_storage:
                        self.metrics_storage[metric_name] = []
                    
                    self.metrics_storage[metric_name].append({
                        'timestamp': timestamp,
                        'value': metric_data['value'],
                        'labels': metric_data['labels']
                    })
                
                # 保持最近1小时的数据
                self.cleanup_old_metrics()
                
                print(f"✅ 成功抓取 {job}({target}) - {len(metrics)} 个指标")
                return True
            else:
                print(f"❌ 抓取失败 {job}({target}) - HTTP {response.status_code}")
                return False
                
        except Exception as e:
            print(f"❌ 抓取异常 {job}({target}) - {str(e)}")
            return False
    
    def parse_metrics(self, metrics_text, job, target):
        """解析Prometheus格式的指标数据"""
        metrics = {}
        lines = metrics_text.strip().split('\n')
        
        for line in lines:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            
            # 解析指标行: metric_name{labels} value timestamp
            match = re.match(r'^([a-zA-Z_:][a-zA-Z0-9_:]*)\s*(\{[^}]*\})?\s+([^\s]+)(\s+\d+)?$', line)
            if match:
                metric_name = match.group(1)
                labels_str = match.group(2) or '{}'
                value_str = match.group(3)
                
                try:
                    value = float(value_str)
                except ValueError:
                    continue
                
                # 解析标签
                labels = {'job': job, 'instance': target}
                if labels_str and labels_str != '{}':
                    labels_str = labels_str[1:-1]  # 移除大括号
                    for label_pair in labels_str.split(','):
                        if '=' in label_pair:
                            key, val = label_pair.split('=', 1)
                            key = key.strip()
                            val = val.strip().strip('"')
                            labels[key] = val
                
                metrics[metric_name] = {
                    'value': value,
                    'labels': labels
                }
        
        return metrics
    
    def cleanup_old_metrics(self):
        """清理1小时前的旧数据"""
        cutoff_time = int(time.time()) - 3600  # 1小时前
        
        for metric_name in list(self.metrics_storage.keys()):
            self.metrics_storage[metric_name] = [
                data for data in self.metrics_storage[metric_name]
                if data['timestamp'] > cutoff_time
            ]
            
            # 如果没有数据了，删除这个指标
            if not self.metrics_storage[metric_name]:
                del self.metrics_storage[metric_name]
    
    def query_metrics(self, query):
        """简单的指标查询"""
        results = []
        
        # 简单匹配指标名称
        for metric_name, data_points in self.metrics_storage.items():
            if query in metric_name:
                if data_points:
                    latest_data = max(data_points, key=lambda x: x['timestamp'])
                    results.append({
                        'metric': latest_data['labels'],
                        'value': [latest_data['timestamp'], str(latest_data['value'])]
                    })
        
        return {
            'status': 'success',
            'data': {
                'resultType': 'vector',
                'result': results
            }
        }
    
    def start_scraping(self):
        """启动指标抓取"""
        def scrape_loop():
            while True:
                for target_info in self.targets:
                    threading.Thread(
                        target=self.scrape_target,
                        args=(target_info,),
                        daemon=True
                    ).start()
                
                time.sleep(self.scrape_interval)
        
        scrape_thread = threading.Thread(target=scrape_loop, daemon=True)
        scrape_thread.start()
        print(f"🔄 开始定期抓取指标，间隔 {self.scrape_interval} 秒")

class PrometheusHandler(BaseHTTPRequestHandler):
    def __init__(self, prometheus_server, *args, **kwargs):
        self.prometheus_server = prometheus_server
        super().__init__(*args, **kwargs)
    
    def do_GET(self):
        parsed_url = urlparse(self.path)
        path = parsed_url.path
        query_params = parse_qs(parsed_url.query)
        
        if path == '/':
            self.serve_dashboard()
        elif path == '/api/v1/query':
            self.handle_query(query_params)
        elif path == '/api/v1/targets':
            self.handle_targets()
        elif path == '/metrics':
            self.handle_prometheus_metrics()
        elif path == '/config':
            self.handle_config()
        else:
            self.send_response(404)
            self.end_headers()
    
    def serve_dashboard(self):
        """提供Web仪表板"""
        self.send_response(200)
        self.send_header('Content-type', 'text/html; charset=utf-8')
        self.end_headers()
        
        # 获取当前指标统计
        total_metrics = len(self.prometheus_server.metrics_storage)
        total_targets = len(self.prometheus_server.targets)
        
        # 获取最近的一些指标数据
        recent_metrics = []
        for metric_name, data_points in list(self.prometheus_server.metrics_storage.items())[:10]:
            if data_points:
                latest = max(data_points, key=lambda x: x['timestamp'])
                recent_metrics.append({
                    'name': metric_name,
                    'value': latest['value'],
                    'timestamp': datetime.fromtimestamp(latest['timestamp']).strftime('%H:%M:%S')
                })
        
        html = f"""
        <!DOCTYPE html>
        <html lang="zh-CN">
        <head>
            <meta charset="UTF-8">
            <meta name="viewport" content="width=device-width, initial-scale=1.0">
            <title>Prometheus 监控系统</title>
            <script src="https://cdn.tailwindcss.com"></script>
            <link href="https://cdn.bootcdn.net/ajax/libs/font-awesome/6.4.0/css/all.min.css" rel="stylesheet">
            <script>
                function refreshData() {{
                    location.reload();
                }}
                
                function queryMetrics() {{
                    const query = document.getElementById('queryInput').value;
                    if (!query) return;
                    
                    fetch(`/api/v1/query?query=${{encodeURIComponent(query)}}`)
                        .then(response => response.json())
                        .then(data => {{
                            document.getElementById('queryResult').innerHTML = 
                                '<pre>' + JSON.stringify(data, null, 2) + '</pre>';
                        }})
                        .catch(error => {{
                            document.getElementById('queryResult').innerHTML = 
                                '<div class="text-red-500">查询失败: ' + error + '</div>';
                        }});
                }}
                
                setInterval(refreshData, 30000); // 30秒自动刷新
            </script>
        </head>
        <body class="bg-gray-100 min-h-screen">
            <div class="container mx-auto px-4 py-8">
                <!-- 头部 -->
                <div class="bg-white rounded-lg shadow-md p-6 mb-6">
                    <div class="flex items-center justify-between">
                        <div>
                            <h1 class="text-3xl font-bold text-gray-800 flex items-center">
                                <i class="fas fa-chart-line text-blue-500 mr-3"></i>
                                Prometheus 监控系统
                            </h1>
                            <p class="text-gray-600 mt-2">基于Python实现的Prometheus监控演示系统</p>
                        </div>
                        <button onclick="refreshData()" class="bg-blue-500 hover:bg-blue-600 text-white px-4 py-2 rounded-lg">
                            <i class="fas fa-sync-alt mr-2"></i>刷新数据
                        </button>
                    </div>
                </div>
                
                <!-- 统计卡片 -->
                <div class="grid grid-cols-1 md:grid-cols-3 gap-6 mb-6">
                    <div class="bg-white rounded-lg shadow-md p-6">
                        <div class="flex items-center">
                            <div class="bg-blue-100 p-3 rounded-full">
                                <i class="fas fa-database text-blue-500 text-xl"></i>
                            </div>
                            <div class="ml-4">
                                <h3 class="text-lg font-semibold text-gray-800">总指标数</h3>
                                <p class="text-2xl font-bold text-blue-500">{total_metrics}</p>
                            </div>
                        </div>
                    </div>
                    
                    <div class="bg-white rounded-lg shadow-md p-6">
                        <div class="flex items-center">
                            <div class="bg-green-100 p-3 rounded-full">
                                <i class="fas fa-bullseye text-green-500 text-xl"></i>
                            </div>
                            <div class="ml-4">
                                <h3 class="text-lg font-semibold text-gray-800">监控目标</h3>
                                <p class="text-2xl font-bold text-green-500">{total_targets}</p>
                            </div>
                        </div>
                    </div>
                    
                    <div class="bg-white rounded-lg shadow-md p-6">
                        <div class="flex items-center">
                            <div class="bg-purple-100 p-3 rounded-full">
                                <i class="fas fa-clock text-purple-500 text-xl"></i>
                            </div>
                            <div class="ml-4">
                                <h3 class="text-lg font-semibold text-gray-800">抓取间隔</h3>
                                <p class="text-2xl font-bold text-purple-500">{self.prometheus_server.scrape_interval}s</p>
                            </div>
                        </div>
                    </div>
                </div>
                
                <!-- 查询界面 -->
                <div class="bg-white rounded-lg shadow-md p-6 mb-6">
                    <h2 class="text-xl font-bold text-gray-800 mb-4">
                        <i class="fas fa-search mr-2"></i>指标查询
                    </h2>
                    <div class="flex gap-4 mb-4">
                        <input type="text" id="queryInput" placeholder="输入指标名称，如: node_cpu" 
                               class="flex-1 px-4 py-2 border border-gray-300 rounded-lg focus:outline-none focus:ring-2 focus:ring-blue-500">
                        <button onclick="queryMetrics()" class="bg-blue-500 hover:bg-blue-600 text-white px-6 py-2 rounded-lg">
                            查询
                        </button>
                    </div>
                    <div id="queryResult" class="bg-gray-50 p-4 rounded-lg min-h-32 overflow-auto">
                        <p class="text-gray-500">在上方输入指标名称进行查询</p>
                    </div>
                </div>
                
                <!-- 最近指标 -->
                <div class="bg-white rounded-lg shadow-md p-6 mb-6">
                    <h2 class="text-xl font-bold text-gray-800 mb-4">
                        <i class="fas fa-list mr-2"></i>最近指标数据
                    </h2>
                    <div class="overflow-x-auto">
                        <table class="w-full table-auto">
                            <thead>
                                <tr class="bg-gray-50">
                                    <th class="px-4 py-2 text-left">指标名称</th>
                                    <th class="px-4 py-2 text-left">当前值</th>
                                    <th class="px-4 py-2 text-left">更新时间</th>
                                </tr>
                            </thead>
                            <tbody>
        """
        
        for metric in recent_metrics:
            html += f"""
                                <tr class="border-t">
                                    <td class="px-4 py-2 font-mono text-sm">{metric['name']}</td>
                                    <td class="px-4 py-2">{metric['value']}</td>
                                    <td class="px-4 py-2">{metric['timestamp']}</td>
                                </tr>
            """
        
        html += """
                            </tbody>
                        </table>
                    </div>
                </div>
                
                <!-- 监控目标 -->
                <div class="bg-white rounded-lg shadow-md p-6">
                    <h2 class="text-xl font-bold text-gray-800 mb-4">
                        <i class="fas fa-bullseye mr-2"></i>监控目标
                    </h2>
                    <div class="grid grid-cols-1 md:grid-cols-2 gap-4">
        """
        
        for target in self.prometheus_server.targets:
            html += f"""
                        <div class="border border-gray-200 rounded-lg p-4">
                            <div class="flex items-center justify-between">
                                <div>
                                    <h3 class="font-semibold text-gray-800">{target['job']}</h3>
                                    <p class="text-gray-600 text-sm">{target['target']}</p>
                                </div>
                                <div class="text-green-500">
                                    <i class="fas fa-check-circle"></i>
                                </div>
                            </div>
                        </div>
            """
        
        html += """
                    </div>
                </div>
                
                <!-- API 链接 -->
                <div class="mt-6 text-center">
                    <div class="bg-gray-50 rounded-lg p-4">
                        <h3 class="text-lg font-semibold text-gray-800 mb-2">API 端点</h3>
                        <div class="flex flex-wrap justify-center gap-4">
                            <a href="/api/v1/targets" class="text-blue-500 hover:text-blue-600">
                                <i class="fas fa-link mr-1"></i>/api/v1/targets
                            </a>
                            <a href="/metrics" class="text-blue-500 hover:text-blue-600">
                                <i class="fas fa-link mr-1"></i>/metrics
                            </a>
                            <a href="/config" class="text-blue-500 hover:text-blue-600">
                                <i class="fas fa-link mr-1"></i>/config
                            </a>
                        </div>
                    </div>
                </div>
            </div>
        </body>
        </html>
        """
        
        self.wfile.write(html.encode('utf-8'))
    
    def handle_query(self, query_params):
        """处理查询请求"""
        self.send_response(200)
        self.send_header('Content-type', 'application/json')
        self.end_headers()
        
        query = query_params.get('query', [''])[0]
        result = self.prometheus_server.query_metrics(query)
        
        self.wfile.write(json.dumps(result).encode('utf-8'))
    
    def handle_targets(self):
        """处理目标列表请求"""
        self.send_response(200)
        self.send_header('Content-type', 'application/json')
        self.end_headers()
        
        targets_info = {
            'status': 'success',
            'data': {
                'activeTargets': [
                    {
                        'discoveredLabels': {'job': target['job']},
                        'labels': {'job': target['job'], 'instance': target['target']},
                        'scrapeUrl': f"http://{target['target']}/metrics",
                        'health': 'up',
                        'lastScrape': datetime.now().isoformat(),
                        'scrapeInterval': target['scrape_interval']
                    }
                    for target in self.prometheus_server.targets
                ]
            }
        }
        
        self.wfile.write(json.dumps(targets_info, indent=2).encode('utf-8'))
    
    def handle_prometheus_metrics(self):
        """处理Prometheus自身的指标"""
        self.send_response(200)
        self.send_header('Content-type', 'text/plain; version=0.0.4; charset=utf-8')
        self.end_headers()
        
        metrics = f"""# HELP prometheus_build_info A metric with a constant '1' value labeled by version, revision, branch, and goversion from which prometheus was built.
# TYPE prometheus_build_info gauge
prometheus_build_info{{version="2.40.0",revision="python-mock"}} 1

# HELP prometheus_config_last_reload_successful Whether the last configuration reload attempt was successful.
# TYPE prometheus_config_last_reload_successful gauge
prometheus_config_last_reload_successful 1

# HELP prometheus_tsdb_symbol_table_size_bytes Size of symbol table in memory.
# TYPE prometheus_tsdb_symbol_table_size_bytes gauge
prometheus_tsdb_symbol_table_size_bytes {len(self.prometheus_server.metrics_storage) * 1000}

# HELP prometheus_tsdb_head_samples_appended_total Total number of appended samples.
# TYPE prometheus_tsdb_head_samples_appended_total counter
prometheus_tsdb_head_samples_appended_total {sum(len(data) for data in self.prometheus_server.metrics_storage.values())}

# HELP up Was the last scrape of the target successful.
# TYPE up gauge
up 1
"""
        
        self.wfile.write(metrics.encode('utf-8'))
    
    def handle_config(self):
        """处理配置查看请求"""
        self.send_response(200)
        self.send_header('Content-type', 'application/json')
        self.end_headers()
        
        config_info = {
            'status': 'success',
            'data': {
                'yaml': yaml.dump(self.prometheus_server.config, default_flow_style=False)
            }
        }
        
        self.wfile.write(json.dumps(config_info, indent=2).encode('utf-8'))
    
    def log_message(self, format, *args):
        """禁用默认日志输出"""
        pass

def create_handler(prometheus_server):
    """创建处理器工厂函数"""
    def handler(*args, **kwargs):
        return PrometheusHandler(prometheus_server, *args, **kwargs)
    return handler

def main():
    print("🚀 启动 Prometheus 服务器...")
    
    # 创建Prometheus服务器实例
    prometheus_server = PrometheusServer()
    
    # 启动指标抓取
    prometheus_server.start_scraping()
    
    # 启动HTTP服务器
    handler = create_handler(prometheus_server)
    server = HTTPServer(('0.0.0.0', 9090), handler)
    
    print("✅ Prometheus 服务器启动成功!")
    print("🌐 Web界面: http://localhost:9090")
    print("📊 指标端点: http://localhost:9090/metrics")
    print("🎯 目标列表: http://localhost:9090/api/v1/targets")
    print("⚙️  配置查看: http://localhost:9090/config")
    
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n👋 Prometheus 服务器已停止")

if __name__ == '__main__':
    main()